# Dependencies:
# pip3 install opencv-python
# pip3 install numpy
# pip3 install scikit-image
import cv2
import os
import numpy as np
from skimage.metrics import structural_similarity

# 测试集目录
original_dir = 'E:\\Datasets\\s2_data\\data\\test'
# 标记篡改位置的图像目录
tamper_dir = 'E:\\Datasets\\s2_data\\data\\test_results'
# output mask directory
dest_dir = 'E:\\Datasets\\s2_data\\data\\images'

walk_dir = os.walk(tamper_dir)

for root, _, file_list in walk_dir:
    for file_name in file_list:
        if not file_name.endswith('.png'):
            continue
        print(file_name)
        orig_path = os.path.join(original_dir, '{}.jpg'.format(file_name[:-4]))
        tamper_path = os.path.join(tamper_dir, file_name)

        orig_img = cv2.imread(orig_path, 1)
        tamper_img = cv2.imread(tamper_path, 1)

        # convert the images to grayscale
        orig_gray = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY)
        tamper_gray = cv2.cvtColor(tamper_img, cv2.COLOR_BGR2GRAY)

        # compute the Structural Similarity Index, return differences
        _, diff = structural_similarity(orig_gray, tamper_gray, full=True)
        diff = (diff * 255).astype("uint8")

        # obtain the different regions
        thresh = cv2.threshold(diff, 0, 255, cv2.THRESH_BINARY_INV)[1]

        # dilate operation
        kernel = np.ones((5, 5), np.uint8)
        dilation = cv2.dilate(thresh, kernel, iterations=4)

        dest_img_name = file_name[:-4] + '.png'
        dest_path = os.path.join(dest_dir, dest_img_name)

        # output the mask
        cv2.imwrite(dest_path, dilation)
